import pyrtl
from pyrtl import *
import re

class Circuit:
    def __init__(self, name=None):
        self.name = name or "circuit"
        self.modules = {}
        self.block = working_block()
        self.wire_signed = {}
        self.reg_list = {}

    def __str__(self):
        return "<Circuit: {}>".format(self.name)

    def add_module(self, name, definition):
        """
        Add a new module to the circuit.

        :param str name: the name of the module to add
        :param Block definition: the Pyrtl Block object that defines the module
        """
        if name in self.modules:
            raise PyrtlError("Module '{}' already exists in this circuit".format(name))
        self.modules[name] = definition

    def get_module(self, name):
        """
        Return the definition for the named module.

        :param str name: the name of the module to retrieve
        :returns: the Block that defines the module
        """
        return self.modules.get(name)

    def remove_module(self, name):
        """
        Remove the named module from the circuit.

        :param str name: the name of the module to remove
        """
        if name not in self.modules:
            raise PyrtlError("Cannot remove nonexistent module '{}'".format(name))
        del self.modules[name]

    def add_wire(self, wirevector):
        """
        Add a new wire to the circuit.

        :param WireVector wirevector: the wire to add to the circuit
        """
        self.block.add_wirevector(wirevector)

    def remove_wire(self, wirevector):
        """
        Remove the given wire from the circuit.

        :param WireVector wirevector: the wire to remove from the circuit
        """
        if wirevector.name not in self.wires:
            raise PyrtlError("Cannot remove nonexistent wire '%s'" % wirevector.name)
        self.block.remove_wirevector(wirevector)

    def add_net(self, logicnet):
        self.block.add_net(logicnet)

    def set_wire_sign(self, wire, signed=False):
        self.wire_signed[wire] = signed

    def get_wire_sign(self, wire):
        return self.wire_signed[wire]

    def get_wire_by_name(self, name):
        return self.block.wirevector_by_name[name]


class FIRRTLParser:
    def __init__(self, filelike, circuit):
        self.filelike = filelike
        self.circuit = circuit
        self.module_name = None
        self.lineno = 0
        self.check_multiple_wires = True
        self.reg_parse = False
        self.cur_reg_name = ''

    def parse(self):
        for line in self.filelike:
            self.lineno += 1
            # print(self.lineno)
            line = line.strip()
            if " @" in line:
                line = line[:line.index(' @')]
            if not line or line.startswith(';'):
                continue
            elif line.startswith('circuit '):
                self._parse_circuit(line)
            elif line.startswith('module '):
                self._parse_module(line)
            elif line.startswith('input'):
                self._parse_input(line)
            elif line.startswith('output'):
                self._parse_output(line)
            elif line.startswith('wire '):
                self._parse_wire(line)
            elif line.startswith('reg '):
                self._parse_register(line)
            elif line.startswith('node '):
                self._parse_node(line)
            elif line.startswith('cmem '):
                self._parse_cmem(line)
            elif line.startswith('smem '):
                self._parse_smem(line)
            elif line.startswith('inst '):
                self._parse_inst(line)
            elif line.startswith('wdef '):
                self._parse_wdef(line)
            else:
                if "<=" in line:
                    self._parse_assign(line)
                elif "=>" in line:
                    self._parse_readport(line)
                elif "mem" in line:
                    continue
                elif "skip" in line:
                    continue
                else:
                    raise PyrtlError("Unknown statement: '{}'".format(line))
    def _parse_readport(self,line):
        if self.reg_parse:
            match_uint = re.match(r'([a-zA-Z0-9_]+)[\s]*=>[\s]*\(UInt<([0-9]+)>\("([a-zA-Z0-9]+)"\),[\s]*([a-zA-Z0-9_]+)\)', line) #reset => (UInt<3>("h0"), a)
            match_sint = re.match(r'([a-zA-Z0-9_]+)[\s]*=>[\s]*\(SInt<([0-9]+)>\("([a-zA-Z0-9]+)"\),[\s]*([a-zA-Z0-9_]+)\)', line)  # reset => (SInt<3>("h0"), a)

            signed = False
            if not match_uint and not match_sint:
                raise PyrtlError("Invalid register statement on line {}: {}".format(self.lineno, line))

            elif match_uint:
                reset = self.circuit.block.wirevector_by_name[match_uint.group(1)]
                bitwidth = int(match_uint.group(2))
                const_val_str = match_uint.group(3)
                wire_name = match_uint.group(4) + "_stupidin"
            elif match_sint:
                reset = self.circuit.block.wirevector_by_name[match_sint.group(1)]
                bitwidth = int(match_sint.group(2))
                const_val_str = match_sint.group(3)
                wire_name = match_sint.group(4) + "_stupidin"
                signed = True

            reg = self.circuit.block.wirevector_by_name[self.cur_reg_name]

            if const_val_str.startswith('h'):
                const_val = int(const_val_str[1:],16)
            elif const_val_str.startswith('b'):
                const_val = int(const_val_str[1:], 2)
            elif const_val_str.startswith('d'):
                const_val = int(const_val_str[1:])
            elif const_val_str[0].isdigit():
                const_val = int(const_val_str)
            else:
                raise PyrtlError("常量数值不合法".format(self.lineno, line))
            if wire_name == self.cur_reg_name:
                raise PyrtlError("read port name can not same with register".format(self.lineno, line))

            reset_const = pyrtl.Const(int(const_val), bitwidth=bitwidth, signed=signed, block=self.circuit.block)
            reset_wire = WireVector(bitwidth=bitwidth, name=wire_name, block=self.circuit.block)
            self.circuit.add_wire(reset_const)
            self.circuit.add_wire(reset_wire)
            self.circuit.set_wire_sign(reset_const,signed=signed)
            self.circuit.set_wire_sign(reset_wire,signed=signed)
            self.reg_parse = False
            tmp = select(reset,0,reset_wire)
            reg.next <<= tmp
            self.circuit.reg_list[reg.name] = reset_wire


    def _parse_circuit(self, line):
        match = re.match(r'circuit[\s]+[a-zA-Z0-9_]+[\s]+\:', line)
        if not match:
            raise PyrtlError("Invalid circuit statement on line {}: {}".format(self.lineno, line))
        self.circuit.name = match.group().split(" ")[1]

    def _parse_module(self, line):
        match = re.match(r'module[\s]+[a-zA-Z0-9_]+[\s]+\:', line)
        if not match:
            raise PyrtlError("Invalid module statement on line {}: {}".format(self.lineno, line))
        self.module_name = match.group().split(" ")[1]

    def _parse_input(self, line):
        match_uint = re.match(r'input[\s]+[a-zA-Z0-9_]+[\s]+\:[\s]+(UInt<[0-9]+>|Clock)', line)
        match_sint = re.match(r'input[\s]+[a-zA-Z0-9_]+[\s]+\:[\s]+(SInt<[0-9]+>)', line)

        if match_uint:
            m = match_uint.group().split(" ")
            #print(m)
            if m[3] == "Clock":
                ipt = Const(val=1,name=m[1])
                self.circuit.add_wire(ipt)
                self.circuit.set_wire_sign(ipt)
            else:
                name = m[1]
                bitwidth = int(re.search("\d+", m[3]).group())
                ipt = Input(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
                self.circuit.add_wire(ipt)
                self.circuit.set_wire_sign(ipt)

        elif match_sint:
            m = match_sint.group().split(" ")
            name = m[1]
            bitwidth = int(re.search("\d+", m[3]).group())
            ipt = Input(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
            self.circuit.add_wire(ipt)
            self.circuit.set_wire_sign(ipt, signed=True)
        else:
            raise PyrtlError("Invalid wire statement on line {}: {}".format(self.lineno, line))

    def _parse_output(self, line):
        match_uint = re.match(r'output[\s]+[a-zA-Z0-9_]+[\s]+\:[\s]+(UInt|Clock)<[0-9]+>', line)
        match_sint = re.match(r'output[\s]+[a-zA-Z0-9_]+[\s]+\:[\s]+(SInt|Clock)<[0-9]+>', line)

        if match_uint:
            m = match_uint.group().split(" ")
            name = m[1]
            bitwidth = int(re.search("\d+", m[3]).group())
            opt = WireVector(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
            self.circuit.add_wire(opt)
            self.circuit.set_wire_sign(opt)
        elif match_sint:
            m = match_sint.group().split(" ")
            name = m[1]
            bitwidth = int(re.search("\d+", m[3]).group())
            opt = WireVector(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
            self.circuit.add_wire(opt)
            self.circuit.set_wire_sign(opt, signed=True)
        else:
            raise PyrtlError("Invalid wire statement on line {}: {}".format(self.lineno, line))

    def _parse_wire(self, line):
        match_uint = re.match(r'wire[\s]+[a-zA-Z0-9_]+[\s]+\:[\s]+(UInt<[0-9]+>|Clock)', line)
        match_sint = re.match(r'wire[\s]+[a-zA-Z0-9_]+[\s]+\:[\s]+(SInt<[0-9]+>)', line)
        
        if match_uint:
            m = match_uint.group().split(" ")
            if m[3] == "Clock":
                ipt = WireVector(bitwidth=1, name=m[1], block=self.circuit.block)
                self.circuit.add_wire(ipt)
                self.circuit.set_wire_sign(ipt)
            else:
                name = m[1]
                bitwidth = int(re.search("\d+", m[3]).group())
                wire = WireVector(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
                self.circuit.add_wire(wire)
                self.circuit.set_wire_sign(wire)

        elif match_sint:
            m = match_sint.group().split(" ")
            name = m[1]
            bitwidth = int(re.search("\d+", m[3]).group())
            wire = WireVector(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
            self.circuit.add_wire(wire)
            self.circuit.set_wire_sign(wire, signed=True)
        else:
            raise PyrtlError("Invalid wire statement on line {}: {}".format(self.lineno, line))

    def _parse_register(self, line):
        match_uint = re.match(r'reg[\s]+([a-zA-Z0-9_]+)[\s]+\:[\s]+(UInt<[0-9]+>)[\s]*,\s*([a-zA-Z0-9_]+)[\s]with[\s]\:', line)
        match_sint = re.match(r'reg[\s]+([a-zA-Z0-9_]+)[\s]+\:[\s]+(SInt<[0-9]+>)[\s]*,\s*([a-zA-Z0-9_]+)[\s]with[\s]\:', line)
        signed = False
        m = []
        if not match_uint and not match_sint:
            raise PyrtlError("Invalid register statement on line {}: {}".format(self.lineno, line))

        elif match_uint:
            m = match_uint.group().split(" ")
            self.reg_parse = True
        else:
            m = match_sint.group().split(" ")
            signed = True
            self.reg_parse = True
        name = m[1]
        self.cur_reg_name = name
        bitwidth = int(re.search("\d+", m[3]).group())
        reg = Register(bitwidth=int(bitwidth), name=name, block=self.circuit.block)
        self.circuit.add_wire(reg)
        self.circuit.set_wire_sign(reg, signed=signed)

        # if self.check_multiple_wires and name in self.circuit.wires:
        #     raise PyrtlError("Multiple wires with the name '{}'".format(name))
        # self.circuit.add_module(
        #     name,
        #     Register(bitwidth=int(bitwidth), name=name))

    def _parse_node(self, line):
        match = re.match(r'node\s+(\S+)\s*=\s*(.*)', line)
        if not match:
            raise PyrtlError("Invalid node statement on line {}: {}".format(self.lineno, line))
        name, expr_text = match.groups()
        # print(name)
        # print(expr_text)
        arg, resultlen, signed = self.parse_expr(expr_text)
        # print(arg)
        # print(resultlen)
        # print(signed)
        dest = WireVector(bitwidth=resultlen, name=name, block=self.circuit.block)
        # print(dest)
        self.circuit.add_wire(dest)
        self.circuit.set_wire_sign(dest, signed=signed)
        dest <<= arg
        # self.circuit.add_net(LogicNet(op='w', op_param=None, args=(arg, ), dests=(dest, )))


    def _parse_assign(self, line):
        match = re.match(r'(\w+)\s+<=\s+(.*)', line)
        if not match:
            raise PyrtlError("Invalid node statement on line {}: {}".format(self.lineno, line))
        name, expr_text = match.groups()
        arg, r, _ = self.parse_expr(expr_text)

        # if len(expr) == 0:
        #     raise PyrtlError("Node expression must have at least one input")
        # if self.check_multiple_wires and name in self.circuit.wires:
        #     raise PyrtlError("Multiple wires with the name '{}'".format(name))
        if name in self.circuit.reg_list:
            wire = self.circuit.reg_list[name]
        else:
            wire = self.circuit.block.wirevector_by_name[name]
        wire <<= arg
        # self.circuit.add_net(LogicNet(op='w', op_param=None, args=(arg, ), dests=(wire, )))

    def parse_arg(self, str):
        if str.startswith("UInt"):
            match = re.match(
                r'UInt<([0-9]+)>\("([a-zA-Z0-9]+)"\)', str)  # (UInt<3>("h0"), a)
            bitwidth = int(match.group(1))
            const_val_str = match.group(2)
            if const_val_str.startswith('h'):
                const_val = int(const_val_str[1:],16)
            elif const_val_str.startswith('b'):
                const_val = int(const_val_str[1:], 2)
            elif const_val_str.startswith('d'):
                const_val = int(const_val_str[1:])
            elif const_val_str[0].isdigit():
                const_val = int(const_val_str)
            else:
                raise PyrtlError("常量数值不合法".format(str))
            arg = Const(const_val, bitwidth)
            self.circuit.set_wire_sign(arg, False)
            return arg
        elif str.startswith("SInt"):
            match = re.match(
                r'SInt<([0-9]+)>\("([-a-zA-Z0-9]+)"\)', str)  # (UInt<3>("h0"), a)
            bitwidth = int(match.group(1))
            const_val_str = match.group(2)
            if const_val_str.startswith('h'):
                const_val = int(const_val_str[1:], 16)
            elif const_val_str.startswith('b'):
                const_val = int(const_val_str[1:], 2)
            elif const_val_str.startswith('d'):
                const_val = int(const_val_str[1:])
            elif const_val_str[0].isdigit():
                const_val = int(const_val_str)
            else:
                raise PyrtlError("常量数值不合法".format(str))
            arg = Const(const_val, bitwidth, signed=True)
            self.circuit.set_wire_sign(arg, True)
            return arg
        elif "(" in str and ")" in str:
            tmpwire , _, signed = self.parse_expr(str)
            self.circuit.set_wire_sign(tmpwire, signed)
            return tmpwire
        else:
            return self.circuit.block.wirevector_by_name[str]

    def simplify_expr_text(self, text):
        text = text[:-1]
        while True:
            if text[0] != "(":
                text = text[1:]
            else:
                return text[1:]
        return pyrtl.PyrtlError("Invalid node statement on line {}: {}", self.lineno, text)


    def parse_expr(self, expr_text):
        tmp_wire = 0
        resultlen = 0
        signed = False
        if expr_text.startswith('add'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
                resultlen = max(resultlen, arg.bitwidth)
            resultlen += 1
            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] + expr[1]
            else:
                tmp_wire = signed_add(expr[0], expr[1])
                signed = True

        elif expr_text.startswith('sub'):  # Subtract Operation
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
                resultlen = max(resultlen, arg.bitwidth)
            resultlen += 1
            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] - expr[1]
            else:
                if isinstance(expr[0], int):
                    expr[0] = Const(expr[0], signed=True)
                if isinstance(expr[1], int):
                    expr[1] = Const(expr[1], signed=True)
                    # if not a wirevector yet, use standard conversion method
                expr[0], expr[1] = match_bitwidth(as_wires(expr[0]), as_wires(expr[1]), signed=True)
                ext_0 = expr[0].sign_extended(resultlen)
                ext_1 = expr[1].sign_extended(resultlen)
                tmp_wire = ext_0 - ext_1
                signed = True

        elif expr_text.startswith('mul'):  # Multiply Operation
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] * expr[1]
            else:
                tmp_wire = signed_mult(expr[0], expr[1])
                signed = True
            resultlen = tmp_wire.bitwidth

        # elif expr_text.startswith('div'):  # Divide Operation
        #     expr = []
        #     args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
        #     for i in args:
        #         if i.startswith("UInt"):
        #             pass
        #         elif i.startswith("SInt"):
        #             pass
        #         else:
        #             arg = self.circuit.block.wirevector_by_name[i]
        #             expr.append(arg)
        #
        #     if not self.circuit.get_wire_sign(expr[0]):
        #         # initial
        #         dividend_bits = expr[0].bitwidth
        #         divisor_bits = expr[1].bitwidth
        #
        #         divisor_is_zero = pyrtl.WireVector(1)
        #         divisor_is_zero <<= (expr[1] == 0)
        #
        #         quotient_wire = pyrtl.WireVector(dividend_bits)
        #         remainder_wire = pyrtl.WireVector(divisor_bits)
        #
        #         remainder_reg = pyrtl.Register(bitwidth=divisor_bits)
        #
        #             # Shift remainder left by 1 bit and add current bit of dividend
        #             remainder_wire <<= pyrtl.concat(remainder_reg[1:], expr[0][-i - 1])
        #
        #             # Subtract divisor from remainder if it fits
        #             subtracted = remainder_wire >= expr[1]
        #             remainder_reg.next <<= pyrtl.select(subtracted, remainder_wire - expr[1], remainder_wire)
        #
        #             # Set current bit of quotient to 1 if divisor was subtracted
        #             quotient_wire |= pyrtl.select(subtracted, 1, 0)
        #             shift_left_logical(quotient_wire, 1)
        #
        #             # Set quotient to 0 if divisor is zero
        #         quotient_wire = pyrtl.select(divisor_is_zero, 0, quotient_wire)
        #         tmp_wire = quotient_wire
        #         resultlen = expr[0].bitwidth
        #     else:
        #         if isinstance(expr[0], int):
        #             expr[0] = Const(expr[0], signed=True)
        #         if isinstance(expr[1], int):
        #             expr[1] = Const(expr[1], signed=True)
        #             # if not a wirevector yet, use standard conversion method
        #         expr[0], expr[1] = match_bitwidth(as_wires(expr[0]), as_wires(expr[1]), signed=True)
        #         ext_0 = expr[0].sign_extended(resultlen)
        #         ext_1 = expr[1].sign_extended(resultlen)
        #         tmp_wire = ext_0 / ext_1
        #         signed = True
        #         resultlen = expr[0].bitwidth + 1

        elif expr_text.startswith('pad'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg = self.circuit.block.wirevector_by_name[args[0]]
            n = int(args[1])
            if arg.bitwidth >= n:
                tmp_wire = arg
                resultlen = arg.bitwidth
            else:
                if args[0].startswith("UInt"):
                    arg = self.parse_arg(arg)
                    tmp_wire = arg.zero_extended(arg)
                elif args[0].startswith("SInt"):
                    arg = self.parse_arg(arg)
                    tmp_wire = arg.sign_extended(arg)
                if not self.circuit.get_wire_sign(arg):
                    tmp_wire = arg.zero_extended(n)
                else:
                    tmp_wire = arg.sign_extended(n)
                    signed = True
                resultlen = n

        elif expr_text.startswith('asUInt'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            i = args[0]
            arg = self.parse_arg(i)
            expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0]
            else:
                tmp_wire, = match_bitwidth(as_wires(expr[0]), signed=False)
            resultlen = tmp_wire.bitwidth

        elif expr_text.startswith('asSInt'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            i = args[0]

            arg = self.parse_arg(i)
            expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire, = match_bitwidth(as_wires(expr[0]), signed=True)
            else:
                tmp_wire = expr[0]
            signed = True
            resultlen = tmp_wire.bitwidth

        elif expr_text.startswith('cvt'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            i = args[0]
            arg = self.parse_arg(i)
            expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0].sign_extended(bitwidth=expr[0].bitwidth + 1)
            else:
                tmp_wire = expr[0]
            resultlen = tmp_wire.bitwidth

        elif expr_text.startswith('neg'):  # negate Operation
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            i = args[0]
            arg = self.parse_arg(i)
            expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                expr[0], = match_bitwidth(as_wires(expr[0]), signed=True)
            else:
                pass
            tmp_wire = ~expr[0] + 1
            resultlen = expr[0].bitwidth
            signed = True

        elif expr_text.startswith('shl'):  # negate Operation
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            i = args[0]
            n = int(args[1])

            arg = self.parse_arg(i)
            expr.append(arg)

            ext_bitwidth = expr[0].bitwidth + n
            if not self.circuit.get_wire_sign(expr[0]):
                expr[0] = expr[0].zero_extended(ext_bitwidth)
                tmp_wire = shift_left_logical(expr[0], n)
            else:
                expr[0] = expr[0] = expr[0].sign_extended(ext_bitwidth)
                tmp_wire = shift_left_logical(expr[0], n)
                signed = True
            resultlen = tmp_wire.bitwidth

        elif expr_text.startswith('dshl'):  # negate Operation
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                if i.startswith("UInt"):
                    arg = self.parse_arg(i)
                    expr.append(arg)
                elif i.startswith("SInt"):
                    arg = self.parse_arg(i)
                    expr.append(arg)
                elif i.isdigit():
                    expr.append(int(i))
                else:
                    arg = self.circuit.block.wirevector_by_name[i]
                    expr.append(arg)


            ext_bitwidth = 0
            if isinstance(expr[1], int):
                ext_bitwidth = expr[0].bitwidth + expr[1]
            else:
                ext_bitwidth = expr[0].bitwidth + 2**expr[1].bitwidth

            if not self.circuit.get_wire_sign(expr[0]):
                expr[0] = expr[0].zero_extended(ext_bitwidth)
                tmp_wire = shift_left_logical(expr[0], expr[1])
            else:
                expr[0] = expr[0].sign_extended(ext_bitwidth)
                tmp_wire = shift_left_logical(expr[0], expr[1])
                signed = True
            resultlen = tmp_wire.bitwidth

        elif expr_text.startswith('shr'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            i = args[0]
            n = int(args[1])

            arg = self.parse_arg(i)
            expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = shift_right_logical(expr[0], n)[:(expr[0].bitwidth-n)]
            else:
                tmp_wire = shift_right_logical(expr[0], n)[:(expr[0].bitwidth-n)]
                signed = True
            resultlen = tmp_wire.bitwidth

        elif expr_text.startswith('dshr'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                if i.startswith("UInt"):
                    arg = self.parse_arg(i)
                    expr.append(arg)
                elif i.startswith("SInt"):
                    arg = self.parse_arg(i)
                    expr.append(arg)
                elif i.isdigit():
                    expr.append(int(i))
                else:
                    arg = self.circuit.block.wirevector_by_name[i]
                    expr.append(arg)

            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = shift_left_logical(expr[0], expr[1])
            else:
                tmp_wire = shift_left_logical(expr[0], expr[1])
                signed = True
            resultlen = tmp_wire.bitwidth


        elif expr_text.startswith('tail'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.circuit.block.wirevector_by_name[args[0]]
            resultlen = len(arg_wire) - int(args[1])
            tmp_wire = arg_wire[:-1*int(args[1])]
        elif expr_text.startswith('head'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.circuit.block.wirevector_by_name[args[0]]
            tmp_wire = arg_wire[arg_wire.bitwidth-int(args[1]):]
            resultlen = int(args[1])
        elif expr_text.startswith('bits'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.circuit.block.wirevector_by_name[args[0]]
            resultlen = int(args[1]) - int(args[2]) + 1
            tmp_wire = arg_wire[int(args[2]):(int(args[1])+1)]
        elif expr_text.startswith('cat'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            tmp_wire = concat(expr[0], expr[1])
            resultlen = expr[0].bitwidth + expr[1].bitwidth
        elif expr_text.startswith('lt'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] < expr[1]
            else:
                tmp_wire = signed_lt(expr[0], expr[1])
            resultlen = 1
        elif expr_text.startswith('leq'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] <= expr[1]
            else:
                tmp_wire = signed_le(expr[0], expr[1])
            resultlen = 1
        elif expr_text.startswith('gt'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] > expr[1]
            else:
                tmp_wire = signed_gt(expr[0], expr[1])
            resultlen = 1
        elif expr_text.startswith('geq'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            if not self.circuit.get_wire_sign(expr[0]):
                tmp_wire = expr[0] >= expr[1]
            else:
                tmp_wire = signed_ge(expr[0], expr[1])
            resultlen = 1
        elif expr_text.startswith('eq'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            tmp_wire = expr[0] == expr[1]
            resultlen = 1
        elif expr_text.startswith('neq'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            tmp_wire = expr[0] != expr[1]
            resultlen = 1
        elif expr_text.startswith('not'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.parse_arg(args[0])
            if self.circuit.get_wire_sign(arg_wire):
                signed=True
            tmp_wire = ~ arg_wire
            resultlen = arg_wire.bitwidth
        elif expr_text.startswith('and') and not expr_text.startswith('andr'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            tmp_wire = expr[0] & expr[1]
            resultlen = max(expr[0].bitwidth, int(expr[1].bitwidth))
        elif expr_text.startswith('or') and not expr_text.startswith('orr'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            tmp_wire = expr[0] | expr[1]
            resultlen = max(expr[0].bitwidth, int(expr[1].bitwidth))
        elif expr_text.startswith('xor') and not expr_text.startswith('xorr'):
            expr = []
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            for i in args:
                arg = self.parse_arg(i)
                expr.append(arg)
            tmp_wire = expr[0] ^ expr[1]
            resultlen = max(expr[0].bitwidth, int(expr[1].bitwidth))
        elif expr_text.startswith('andr'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.parse_arg(args[0])
            tmp_wire = and_all_bits(arg_wire)
            resultlen = 1
        elif expr_text.startswith('orr'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.parse_arg(args[0])
            tmp_wire = or_all_bits(arg_wire)
            resultlen = 1
        elif expr_text.startswith('xorr'):
            args = re.sub(r'^\w+\(|\)$', '', expr_text).split(", ")
            arg_wire = self.parse_arg(args[0])
            tmp_wire = xor_all_bits(arg_wire)
            resultlen = 1
        elif expr_text.startswith('mux'):
            p2 = re.compile(r"[(](.*)[)]", re.S)
            a = re.findall(p2, expr_text)
            args = re.split(r",(?![^(]*\))", a[0])
            sel = self.parse_arg(args[0].strip())
            t = self.parse_arg(args[1].strip())
            f = self.parse_arg(args[2].strip())
            if self.circuit.get_wire_sign(t):
                t, f = match_bitwidth(t, f, signed=True)
            else:
                t, f = match_bitwidth(t, f, signed=False)
            tmp_wire = select(sel, t, f)
            resultlen = len(t)
        elif "(" not in expr_text:
            #print(expr_text, self.lineno)
            tmp_wire = self.circuit.get_wire_by_name(expr_text)
            resultlen = len(tmp_wire)
        elif "UInt<" in expr_text or "SInt<" in expr_text:
            tmp_wire = self.parse_arg(expr_text)
            resultlen = len(tmp_wire)
        else:
            raise PyrtlError("Unknown statement: '{}'".format(expr_text))
        return tmp_wire, resultlen, signed

    def _parse_cmem(self, line):
        match = re.match(r'cmem\s+(\S+)\s*\[\s*(\d+)\s*\]\s*\[\s*(\d+)\s*\]\s*', line)
        if not match:
            raise PyrtlError("Invalid cmem statement on line {}: {}".format(self.lineno, line))
        name, width, depth = match.groups()
        if self.check_multiple_wires and name in self.circuit.wires:
            raise PyrtlError("Multiple wires with the name '{}'".format(name))
        self.circuit.add_module(
            name,
            RegisterFile(bitwidth=int(width), write_ports=1, size=int(depth)))

    def _parse_smem(self, line):
        match = re.match(r'smem\s+(\S+)\s*\[\s*(\d+)\s*\]\s*\[\s*(\d+)\s*\]\s*', line)
        if not match:
            raise PyrtlError("Invalid smem statement on line {}: {}".format(self.lineno, line))
        name, width, depth = match.groups()
        if self.check_multiple_wires and name in self.circuit.wires:
            raise PyrtlError("Multiple wires with the name '{}'".format(name))
        self.circuit.add_module(
            name,
            RegisterFile(bitwidth=int(width), write_ports=1, size=int(depth), read_ports=1))

    def _parse_inst(self, line):
        match = re.match(r'inst\s+(\S+)\s*of\s+(\S+)', line)
        if not match:
            raise PyrtlError("Invalid inst statement on line {}: {}".format(self.lineno, line))
        instance_name, module_name = match.groups()
        module_def = self.circuit.get_module(module_name)
        if not module_def:
            raise PyrtlError("Unknown module: {}".format(module_name))
        self.circuit.add_module(instance_name, module_def)

    def _parse_wdef(self, line):
        match = re.match(r'wdef\s+(\S+)\s*\=\s*(.*)', line)
        if not match:
            raise PyrtlError("Invalid wdef statement on line {}: {}".format(self.lineno, line))
        name, expr_text = match.groups()
        expr = parse_expr(expr_text)
        if len(expr) != 1:
            raise PyrtlError("WDef expression must have exactly one output")
        if self.check_multiple_wires and name in self.circuit.wires:
            raise PyrtlError("Multiple wires with the name '{}'".format(name))
        wire = WireVector(bitwidth=expr[0].bitwidth, name=name)
        self.circuit._add_wire(wire)
        LogicNet(op='w', op_param=None, args=expr, dests=[wire])

# str = """circuit adder :
#   module adder :
#     input a : UInt<8>
#     input b : UInt<8>
#     output sum : UInt<8>

#     node temp = add(a, b)

#     sum <= tail(temp, 8)"""

if __name__ == "__main__":
    c = Circuit()

    #f = open("./FFT.lo.fir")
    #f = open("./TLI2C.lo.fir")
    #f = open("./TLSPI.lo.fir")
    f = open("./RocketCore.lo.fir")

    parser = FIRRTLParser(f, c)
    parser.parse()

    max_ = 0
    com_time = 0
    ess_time = 0
    num = 1
    for i in range(num):
        print(i)
        input_width = []
        from pyrtl.wire import Input

        IN = pyrtl.working_block().wirevector_subset(Input)
        for inn in IN:
            name = inn.name
            width = inn.bitwidth
            input_width.append([name, width])
        import random

        inputs_random = {}
        inputs_select = {}
        inputs_old = {}
        inputs_ = {}

        for inp in input_width:
            inputs_random[inp[0]] = []
            inputs_select[inp[0]] = []
            result = random.randint(0, 2 ** inp[1] - 1)
            inputs_old[inp[0]] = result

        for j in range(131072):
            for inp in input_width:
                if inp[1] > 10:
                    inp[1] = 10
                result = random.randint(0, 2 ** inp[1] - 1)
                inputs_[inp[0]] = result
                inputs_random[inp[0]].append(result)
            if random.random() < 0.2:
                for name in inputs_old:
                    inputs_select[name].append(inputs_old[name])
            else:
                for name in inputs_:
                    inputs_select[name].append(inputs_[name])
                    inputs_old[name] = inputs_[name]

    #sim = pyrtl.GPUSim_loopcpugpu_asyncstream_dll.GPUSim('i2c_stream',64,20)
    #sim = pyrtl.GPUSim_now_loopcpugpu_dll.GPUSim('rock_reg', 64, 20)
    sim = pyrtl.GPUSim_now_nonloop.GPUSim('rock', 64, 20)
    sim.make_input(inputs_select)